# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.

import torch
import numpy
import numpy as np
from atk.configs.dataset_config import InputDataset

from atk.tasks.api_execute import register
from atk.tasks.api_execute.base_api import BaseApi

@register("function_aclnn_moe_compute_expert_tokens")
class FunctionApi(BaseApi):
    def __call__(self, input_data: InputDataset, with_output: bool = False):
        sorted_experts = input_data.kwargs.get("sortedExperts").cpu().numpy()
        num_experts = input_data.kwargs.get("numExperts")
        arr_length = sorted_experts.shape[-1]
        res = np.arange(num_experts)
        for i in range(num_experts):
            target = i
            low = 0
            high = arr_length - 1
            target_location = -1
            while low <= high:
                mid = (low + high) // 2
                if sorted_experts[mid] > target:
                    high = mid - 1
                else:
                    low = mid + 1
                    target_location = mid
            res[i] = target_location + 1
        res = res.astype(np.int32)
        return torch.from_numpy(res)
    
    def init_by_input_data(self, input_data: InputDataset):
        sorted_experts = input_data.kwargs["sortedExperts"]

        input_data.kwargs['sortedExperts'] = torch.sort(sorted_experts, stable=True)[0]


    def get_cpp_func_signature_type(self):
        return "aclnnStatus aclnnMoeComputeExpertTokensGetWorkspaceSize(const aclTensor* sortedExperts, int64_t numExperts, const aclTensor* out, uint64_t *workspaceSize, aclOpExecutor **executor)"